import torch

# 加载TorchScript模型
model = torch.jit.load("data/weights/yolopv2 copy 2.pt")
model.eval()  # 设置为评估模式

# 创建示例输入（根据模型输入尺寸调整）
input_shape = (1, 3, 640, 640)  # 示例输入尺寸：(batch, channels, height, width)
dummy_input = torch.randn(*input_shape)

# 导出为ONNX
onnx_path = "model.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    opset_version=12,  # 根据需求调整opset版本
    input_names=["input"],
    output_names=["output"],
)

# 使用Netron打开ONNX模型
import netron
netron.start(onnx_path)